import sys, os, copy, math
from topology_engineer import *
sys.path.append("..")
#import aurora_network
from gurobipy import *

'''		
Design fractional topology based on 
traffic clusters using LP formulation - Permits only direct hops, so traffic MUST traverse directly from one pod to another
'''
class DirectPathTopologyEngineer(TopologyEngineer):
	def __init__(self, aurora_network, reconfig_length, training_length, numK=1, return_predicted_mlu = False):
		TopologyEngineer.__init__(self, aurora_network, 0)
		assert(numK >= 1)
		self.numK = numK
		self.training_length = training_length
		self.reconfig_length = reconfig_length
		self.return_predicted_mlu = return_predicted_mlu
		return

	def get_filename_param(self):
		return "direct_toe_r{}t{}k{}".format(self.reconfig_length, self.training_length, self.numK)

	# returns a string description of this class
	def get_string(self):
		return "Direct path ToE"

	def is_static(self):
		return False
	
	## done - ish
	# first step of topology design which solves the LP that maximizes traffic scale up
	def _maximize_traffic_scaleup(self, nblocks, traffic_matrices):
		model = Model("maximize traffic scale up direct paths")
		model.setParam( 'OutputFlag', False )
		mu = model.addVar(lb=0, ub=GRB.INFINITY, obj=1., vtype=GRB.CONTINUOUS, name="mu")
		fractional_topology_var = [None] * nblocks
		num_tm = len(traffic_matrices)
		for i in range(nblocks):
			fractional_topology_var[i] = [None] * nblocks
			for j in range(nblocks):
				if i != j:
					upper_bound = min(self.aurora_network.get_num_links(i), self.aurora_network.get_num_links(j))
					fractional_topology_var[i][j] = model.addVar(lb=0, ub=upper_bound, obj=0, vtype=GRB.CONTINUOUS, name="lc" + str(i) + ":" + str(j))

		for i in range(nblocks):
			egress_link_sum = LinExpr()
			ingress_link_sum = LinExpr()
			for j in range(nblocks):
				if i != j:
					egress_link_sum += fractional_topology_var[i][j]
					ingress_link_sum += fractional_topology_var[j][i]
			model.addConstr(lhs=egress_link_sum, sense=GRB.LESS_EQUAL, rhs=self.aurora_network.get_num_links(i))
			model.addConstr(lhs=ingress_link_sum, sense=GRB.LESS_EQUAL, rhs=self.aurora_network.get_num_links(i))

		for tm_index, tm in zip(range(num_tm), traffic_matrices):
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						link_capacity = min(self.aurora_network.get_link_capacity(i), self.aurora_network.get_link_capacity(j))
						model.addConstr(lhs=mu * tm[i][j], sense=GRB.LESS_EQUAL, rhs=fractional_topology_var[i][j] * link_capacity)

		# set up the objective function
		model.setObjective(mu, GRB.MAXIMIZE)
		# start optimizing
		try: 
			model.optimize()
			adj_matrix = np.zeros((nblocks, nblocks))
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						adj_matrix[i][j] = fractional_topology_var[i][j].x
			return adj_matrix, mu.x
		except GurobiError as e:
			print ("Maximize Traffic Scale Up: Error code " + str(e. errno ) + ": " + str(e))
			return None
		except AttributeError :
			print ("Maximize Traffic Scale Up: Encountered an attribute error ")
			return None

	# conducts topology engineering given the representative traffic matrices
	# aurora network - Aurora Network, contains all the parameters
	# training_traffic_snapshots - all the snapshots used for training the topology
	# number of k - which is the number of traffic clusters we use to solve engineer topology
	def topology_engineer_given_representative_TMs(self, traffic_matrices, all_paths):
		# first get number of blocks and number of links and link capacity
		nblocks = self.aurora_network.get_num_blocks()
		adj_matrix_spare, scaleup = self._maximize_traffic_scaleup(nblocks, traffic_matrices)
		#print(adj_matrix_spare)
		adj_matrix = self.progressive_filling(adj_matrix_spare)
		if adj_matrix is None:
			adj_matrix = adj_matrix_spare
		if self.return_predicted_mlu:
			return adj_matrix, None, 1./scaleup
		else:
			return adj_matrix_spare, None